import os
import subprocess
import re
import logging

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def get_gpu_memory_map():
    """
    获取GPU内存使用情况映射
    
    Returns:
        dict: GPU ID到可用内存（MB）的映射
    """
    try:
        # 使用nvidia-smi命令获取GPU信息
        if os.system('which nvidia-smi > /dev/null 2>&1') == 0:
            result = subprocess.check_output(['nvidia-smi', '--query-gpu=index,memory.free', '--format=csv,nounits,noheader'])
            result = result.decode('utf-8').strip()
            
            gpu_memory_map = {}
            for line in result.split('\n'):
                if line.strip():
                    gpu_id, memory_free = line.split(',')
                    gpu_memory_map[int(gpu_id)] = int(memory_free.strip())
            
            return gpu_memory_map
        else:
            logger.warning("nvidia-smi命令不可用，无法获取GPU内存信息")
            return {}
    except Exception as e:
        logger.error(f"获取GPU内存信息时出错: {str(e)}")
        return {}